package com.twb.auth.dataScope.interceptor;

import java.lang.reflect.Method;
import java.sql.SQLException;
import java.util.List;
import java.util.Properties;
import java.util.Set;

import javax.servlet.http.HttpServletRequest;

import org.apache.commons.collections4.CollectionUtils;
import org.apache.commons.lang3.StringUtils;
import org.apache.ibatis.cache.CacheKey;
import org.apache.ibatis.executor.Executor;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.mapping.SqlCommandType;
import org.apache.ibatis.plugin.Interceptor;
import org.apache.ibatis.plugin.Intercepts;
import org.apache.ibatis.plugin.Invocation;
import org.apache.ibatis.plugin.Plugin;
import org.apache.ibatis.plugin.Signature;
import org.apache.ibatis.session.ResultHandler;
import org.apache.ibatis.session.RowBounds;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.stereotype.Component;
import org.springframework.web.context.request.RequestContextHolder;

import com.twb.auth.dataScope.annotation.DataScope;
import com.twb.auth.dataScope.enums.DataScopeType;
import com.twb.core.base.entity.BaseLoginUser;
import com.twb.core.config.TwbConfigProperties;
import com.twb.core.security.TwbSecurityUtil;

import net.sf.jsqlparser.expression.Alias;
import net.sf.jsqlparser.expression.Expression;
import net.sf.jsqlparser.expression.Parenthesis;
import net.sf.jsqlparser.expression.StringValue;
import net.sf.jsqlparser.expression.operators.conditional.AndExpression;
import net.sf.jsqlparser.expression.operators.conditional.OrExpression;
import net.sf.jsqlparser.expression.operators.relational.EqualsTo;
import net.sf.jsqlparser.expression.operators.relational.ExpressionList;
import net.sf.jsqlparser.expression.operators.relational.InExpression;
import net.sf.jsqlparser.parser.CCJSqlParserUtil;
import net.sf.jsqlparser.schema.Column;
import net.sf.jsqlparser.schema.Table;
import net.sf.jsqlparser.statement.Statement;
import net.sf.jsqlparser.statement.select.FromItem;
import net.sf.jsqlparser.statement.select.PlainSelect;
import net.sf.jsqlparser.statement.select.Select;
import net.sf.jsqlparser.statement.select.SelectBody;
import net.sf.jsqlparser.statement.select.SetOperationList;
import net.sf.jsqlparser.statement.select.SubSelect;
import net.sf.jsqlparser.util.TablesNamesFinder;

/**
 * @ClassName DataScopeInterceptor
 * @Description 数据权限拦截器，数据类型的排序是：全部数据> 本公司及以下数据 > 本公司数据 > 本部门及以下数据 > 本部门数据 > 仅本人数据
 * @author DSY
 * @date 2022年7月26日 上午9:06:46
 */
@Component
@ConditionalOnProperty(prefix = "twb.dataScope", name = "enabled", havingValue = "true")
@Intercepts({
		@Signature(type = Executor.class, method = "query", args = { MappedStatement.class, Object.class,
				RowBounds.class, ResultHandler.class }),
		@Signature(type = Executor.class, method = "query", args = { MappedStatement.class, Object.class,
				RowBounds.class, ResultHandler.class, CacheKey.class, BoundSql.class }), })
public class DataScopeInterceptor implements Interceptor {
	private static final Logger log = LoggerFactory.getLogger(DataScopeInterceptor.class);

	private static final String CREATOR_ORGID = "creator_orgid";
	private static final String CREATOR_ID = "creator_id";
	@Autowired
	protected TwbSecurityUtil twbSecurityUtil;
	@Autowired
	HttpServletRequest request;
	@Autowired
	protected TwbConfigProperties twbConfigProperties;

	/**
	 * @Title getLoginUserDataScopes
	 * @Description 获取用户的数据权限信息
	 * @param loginUser
	 * @return
	 * @AuthorOriginally DSY
	 * @date 2022年7月25日 下午10:05:22
	 */
	private Set<String> getLoginUserDataScopes(BaseLoginUser loginUser) {
		Set<String> dataScopes = twbSecurityUtil.getLoginUserDataScopes(loginUser.getToken());
		return dataScopes;
	}

	/**
	 * intercept, 此处实现拦截逻辑
	 */
	@Override
	public Object intercept(Invocation invocation) throws Throwable {
		long start = System.currentTimeMillis();
		log.debug("数据权限拦截器开始");
		Object[] args = invocation.getArgs();
		// 取Mapper文件定义
		MappedStatement mappedStatement = (MappedStatement) args[0];
        Object parameter = args[1];
        RowBounds rowBounds = (RowBounds) args[2];
        ResultHandler<?> resultHandler = (ResultHandler<?>) args[3];
        Executor executor = (Executor) invocation.getTarget();
        CacheKey cacheKey;
        BoundSql boundSql;
        //由于逻辑关系，只会进入一次
        if(args.length == 4){
            //4 个参数时
            boundSql = mappedStatement.getBoundSql(parameter);
            cacheKey = executor.createCacheKey(mappedStatement, parameter, rowBounds, boundSql);
        } else {
            //6 个参数时
            cacheKey = (CacheKey) args[4];
            boundSql = (BoundSql) args[5];
        }

		// 取绑定的SQL脚本并打印
		String sql = boundSql.getSql();
		log.debug("Intercept SQL: {}", sql);

		String delegateSql = sql;
		// 获取执行方法的位置
		String namespace = mappedStatement.getId();
		// 获取mapper名称
		String className = namespace.substring(0, namespace.lastIndexOf("."));
		// 获取方法名
		String methodName = namespace.substring(namespace.lastIndexOf(".") + 1, namespace.length());
		methodName = methodName.replace("_COUNT", "");
		// 获取当前mapper 的方法
		Method[] ms = Class.forName(className).getMethods();
		Boolean ignore = true;
		DataScope dataScope = null;
		for (Method m : ms) {
			if (m.getName().equals(methodName)) {
				// 获取注解 来判断是不是要拦截sql
				dataScope = m.getAnnotation(DataScope.class);
				if (dataScope != null && !dataScope.ignore()) {
					ignore = false;
				}
				break;
			}
		}
		if (ignore) {
			log.debug("数据权限拦截器未拦截");
			return invocation.proceed();
		}
		// 如果不是http请求的查询则不进行拦截
		if (RequestContextHolder.getRequestAttributes() != null) {
			BaseLoginUser loginUser = this.twbSecurityUtil.getLoginUser(this.request);
			// 不是select方法、表名不符合规则不进行拦截
			if (loginUser == null || !(SqlCommandType.SELECT.equals(mappedStatement.getSqlCommandType()))
					|| StringUtils.isBlank(delegateSql) || loginUser.isDeveloper() || loginUser.getDataScopeType() == null) {
				return invocation.proceed();
			}
			// 如果是查询全部数据，则直接跳过
			if (loginUser.getDataScopeType().equals(DataScopeType.ALL.getCode())) {
				return invocation.proceed();
			}
			// 核心注入自定义脚本的方法, 具体往下看
			delegateSql = this.delegate(sql, loginUser,dataScope);
			BoundSql newBoundSql = new BoundSql(mappedStatement.getConfiguration(),delegateSql, boundSql.getParameterMappings(), boundSql.getParameterObject());
			log.debug("Delegate SQL: {}", delegateSql);
			long end = System.currentTimeMillis();
			log.debug("数据权限拦截器截止，耗时：" + (end - start));
			return executor.query(mappedStatement, parameter, rowBounds, resultHandler, cacheKey, newBoundSql);
		}
		return invocation.proceed();
	}

	private String delegate(String originSql, BaseLoginUser loginUser,DataScope dataScope) throws Exception {
		return this.rewrite(originSql, loginUser,dataScope);
	}

	/**
	 * SQL重写路由
	 * 
	 * @param originSql
	 * @return
	 * @throws Exception
	 */
	private String rewrite(String originSql, BaseLoginUser loginUser,DataScope dataScope) throws Exception {
		Statement statement = CCJSqlParserUtil.parse(originSql);
		if (statement instanceof Select) {
			return this.rewriteSelectSql(statement, loginUser,dataScope);
		} else {
			// 自行实现异常
			throw new SQLException();
		}
	}

	private String rewriteSelectSql(Statement statement, BaseLoginUser loginUser,DataScope dataScope) {
		Select selectStatement = (Select) statement;

		TablesNamesFinder tablesNameFinder = new TablesNamesFinder();
		List<String> tableNames = tablesNameFinder.getTableList(selectStatement);
		// select 1 OR select now()
		if (tableNames.size() == 0) {
			return selectStatement.toString();
		}

		// 复杂查询, 譬如JOIN, 普通连表等
		PlainSelect plainSelect = (PlainSelect) selectStatement.getSelectBody();
		FromItem fromItem = plainSelect.getFromItem();
		if (fromItem instanceof SubSelect) {
			SubSelect subSelect = (SubSelect) fromItem;
			if (subSelect.getSelectBody() != null) {
				if (subSelect.getSelectBody() instanceof SetOperationList) {
					SetOperationList operationList = (SetOperationList) subSelect.getSelectBody();
					if (CollectionUtils.isNotEmpty(operationList.getSelects())) {
						List<SelectBody> plainSelects = operationList.getSelects();
						for (SelectBody selectBody : plainSelects) {
							PlainSelect select = (PlainSelect) selectBody;
							String mainTableName = dealPlainSelect(select);
							setWhere(select, selectStatement, mainTableName, loginUser,dataScope);
						}
					}
				} else {
					PlainSelect select = (PlainSelect) subSelect.getSelectBody();
					String mainTableName = dealPlainSelect(select);
					setWhere(select, selectStatement, mainTableName, loginUser,dataScope);
				}
			}
		} else if (fromItem instanceof Table) {
			String mainTableName = ((Table) plainSelect.getFromItem()).getName();
			setWhere(plainSelect, selectStatement, mainTableName, loginUser,dataScope);
		}
		return selectStatement.toString();
	}

	/**
	 * 多个子查询处理获取主表别名
	 * 
	 * @param plainSelect
	 * @return
	 */
	private String dealPlainSelect(PlainSelect plainSelect) {
		String mainTableName = "";
		FromItem fromItem = plainSelect.getFromItem();
		if (fromItem instanceof Table) {
			mainTableName = ((Table) plainSelect.getFromItem()).getName();
			return mainTableName;
		} else {
			SubSelect subSelect = (SubSelect) fromItem;
			return dealPlainSelect((PlainSelect) subSelect.getSelectBody());
		}
	}

	/**
	 * @Title setWhere
	 * @Description 拼接查询条件
	 * @param plainSelect
	 * @param statement
	 * @param tableName
	 * @param loginUser
	 * @AuthorOriginally DSY
	 * @date 2022年7月25日 下午6:21:19
	 */
	private void setWhere(PlainSelect plainSelect, Statement statement, String tableName, BaseLoginUser loginUser,DataScope dataScope) {
		String aliasName = this.getTableAlias(statement, tableName);
		Set<String> dataScopeIds = getLoginUserDataScopes(loginUser);
		String dataScopeType = loginUser.getDataScopeType();
		String deptAlias = dataScope.deptAlias();//自定义部门表前缀
		String deptColumn = dataScope.deptColumn();//自定义部门表列名
		String userAlias = dataScope.userAlias();//自定义用户表前缀
		String userColumn = dataScope.userColumn();//自定义用户表列名
		if(StringUtils.isBlank(deptColumn)) {
			deptColumn = CREATOR_ORGID;
		}
		if(StringUtils.isBlank(userColumn)) {
			userColumn = CREATOR_ID;
		}
		// 数据类型的排序是：全部数据 > 本公司及以下数据 > 本公司数据> 本部门及以下数据 > 本部门数据 > 仅本人数据
		if (dataScopeType.equals(DataScopeType.DEPT.getCode()) && dataScopeIds.isEmpty()) {
			if(StringUtils.isNotBlank(deptAlias)) {
				aliasName = deptAlias;
			}
			plainSelect.setWhere(this.newAndExpression(aliasName, plainSelect.getWhere(), deptColumn, loginUser.getOrgId()));
		} else if (dataScopeType.equals(DataScopeType.SELF.getCode()) && dataScopeIds.isEmpty()) {
			if(StringUtils.isNotBlank(userAlias)) {
				aliasName = userAlias;
			}
			plainSelect.setWhere(this.newAndExpression(aliasName, plainSelect.getWhere(), userColumn, loginUser.getId()));
		} else if (dataScopeType.equals(DataScopeType.SELF.getCode()) && !dataScopeIds.isEmpty()) {
			String userAliasName = aliasName;
			String deptAliasName = aliasName;
			if(StringUtils.isNotBlank(deptAlias)) {
				deptAliasName = deptAlias;
			}
			if(StringUtils.isNotBlank(userAlias)) {
				userAliasName = userAlias;
			}
			AndExpression andExpression = new AndExpression();
			OrExpression orExpression1 = new OrExpression(
					this.newAndExpression(userAliasName, null, userColumn, loginUser.getId()),
					this.newAndInExpression(deptAliasName, null, deptColumn, dataScopeIds));
			andExpression.setLeftExpression(new Parenthesis(orExpression1));
			andExpression.setRightExpression(new Parenthesis(plainSelect.getWhere()));
			plainSelect.setWhere(andExpression);
		} else if (dataScopeType.equals(DataScopeType.DEPT_WITH_CHILD.getCode())) {
			if(StringUtils.isNotBlank(deptAlias)) {
				aliasName = deptAlias;
			}
			plainSelect.setWhere(this.newAndInExpression(aliasName, plainSelect.getWhere(), deptColumn, dataScopeIds));
		} else {
			if(StringUtils.isNotBlank(deptAlias)) {
				aliasName = deptAlias;
			}
			plainSelect.setWhere(this.newAndInExpression(aliasName, plainSelect.getWhere(), deptColumn, dataScopeIds));
		}
	}

	private Expression newAndInExpression(String aliasName, Expression whereExpression, String column,
			Set<String> values) {
		InExpression newIn = this.newIn(aliasName, column, values);
		if (whereExpression == null) {
			return newIn;
		}
		return new AndExpression(newIn, new Parenthesis(whereExpression));
	}

	private Expression newAndExpression(String aliasName, Expression whereExpression, String column, String value) {
		EqualsTo equalsTo = this.newEqualTo(aliasName, column, value);
		if (whereExpression == null) {
			return equalsTo;
		}
		// rewrite parent where expression
		return new AndExpression(equalsTo, new Parenthesis(whereExpression));
	}

	private InExpression newIn(String aliasName, String column, Set<String> values) {
		InExpression inExpression = new InExpression();
		inExpression.setLeftExpression(new Column((aliasName == null ? "" : aliasName + '.') + column));
		ExpressionList list = new ExpressionList();
		for (String value : values) {
			Expression elements = new StringValue(value);
			list.addExpressions(elements);
		}
		inExpression.setRightItemsList(list);
		return inExpression;
	}

	private EqualsTo newEqualTo(String aliasName, String column, String value) {
		EqualsTo equalsTo = new EqualsTo();
		equalsTo.setLeftExpression(new Column((aliasName == null ? "" : aliasName + '.') + column));
		equalsTo.setRightExpression(new StringValue(value));
		return equalsTo;
	}

	private String getTableAlias(Statement stmt, String tableName) {
		String aliasName = null;
		if (stmt instanceof Select) {
			Select select = (Select) stmt;
			PlainSelect plainSelect = (PlainSelect) select.getSelectBody();
			FromItem fromItem = plainSelect.getFromItem();
			if (fromItem instanceof SubSelect) {
				PlainSelect sel = (PlainSelect) ((SubSelect) fromItem).getSelectBody();
				String tabName = dealPlainSelect(sel);
				if ((tabName.equalsIgnoreCase(tableName))) {
					Alias alias = sel.getFromItem().getAlias();
					aliasName = alias != null ? alias.getName() : null;
				}
			} else if (fromItem instanceof Table) {
				if (((Table) fromItem).getName().equalsIgnoreCase(tableName)) {
					Alias alias = fromItem.getAlias();
					aliasName = alias != null ? alias.getName() : null;
				}
			}

		}
		return aliasName;
	}

	/**
	 * plugin方法是拦截器用于封装目标对象的，通过该方法我们可以返回目标对象本身，也可以返回一个它的代理
	 *
	 * 当返回的是代理的时候我们可以对其中的方法进行拦截来调用intercept方法 -- Plugin.wrap(target, this)
	 * 当返回的是当前对象的时候 就不会调用intercept方法，相当于当前拦截器无效
	 */
	@Override
	public Object plugin(Object target) {
		return Plugin.wrap(target, this);
	}

	@Override
	public void setProperties(Properties properties) {
	}
}
